#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
 
从检测结果生成xml文件

"""

import sys
import argparse
import os
import cv2
from xml.dom.minidom import Document


def is_invalid_big_box(box_bnd, box_rects):
    invalid = False
    in_count = 0
    for box in box_rects:
        if box[0] == box_bnd[0]:
            # 判断box是否在box_bnd内部
            if (box_bnd[1] < box[1]) \
                    and (box_bnd[2] < box[2]) \
                    and (box_bnd[3] > box[3]) \
                    and (box_bnd[4] > box[4]):
                in_count += 1
                if in_count > 3:
                    invalid = True
                    break
    return invalid


def add_object(doc, annotation, obj_name,
               xmin_v, ymin_v, xmax_v, ymax_v, occl):
    obj=doc.createElement("object")
    
    obj_name = doc.createElement("name")
    obj_name.appendChild(doc.createTextNode(obj_name))
    pose = doc.createElement("pose")
    pose.appendChild(doc.createTextNode("front"))  
    truncated = doc.createElement("truncated")
    truncated.appendChild(doc.createTextNode(str(occl)))
    difficult = doc.createElement("difficult")
    difficult.appendChild(doc.createTextNode("0"))
    obj.appendChild(obj_name)
    obj.appendChild(pose)
    obj.appendChild(truncated)
    obj.appendChild(difficult)
    
    bndbox = doc.createElement("bndbox")
    xmin = doc.createElement("xmin")
    ymin = doc.createElement("ymin")
    xmax = doc.createElement("xmax")
    ymax = doc.createElement("ymax")
    xmin.appendChild(doc.createTextNode(str(xmin_v)))
    ymin.appendChild(doc.createTextNode(str(ymin_v)))
    xmax.appendChild(doc.createTextNode(str(xmax_v)))
    ymax.appendChild(doc.createTextNode(str(ymax_v)))
    bndbox.appendChild(xmin)
    bndbox.appendChild(ymin)
    bndbox.appendChild(xmax)
    bndbox.appendChild(ymax)
    obj.appendChild(bndbox)
    
    annotation.appendChild(obj)


def generate_dict_from_detection(detection_result_dir):
    ground_truth_dict = {}
    for root, dirs, files in os.walk(detection_result_dir):
        for file in files:
            fl = os.path.join(detection_result_dir, file)
            if not os.path.exists(fl):
                print(fl, "not exist!")
                continue
            img_name = "{}.jpg".format(file.split('.')[0])
            ground_truth_dict[img_name] = []
            ground_truth = open(fl, 'r')
            if not ground_truth:
                continue
            for line in ground_truth.readlines():
                info = line.strip().split()
                rect = list(map(int, info[2:]))
                rect.insert(0, info[0])
                # info = list(map(int, info))
                #[cls xmin,ymin,xmax,ymax]
                ground_truth_dict[img_name].append(rect)
            ground_truth.close()
    return ground_truth_dict


def detection_to_voc_xml(detection_result_dir, voc_root, labels):
    img_path = os.path.join(voc_root, 'JPEGImages/')
    save_xml_path = os.path.join(voc_root, 'Annotations/')
    if not os.path.exists(save_xml_path):
        os.makedirs(save_xml_path)
    detection_result_num = len(os.listdir(detection_result_dir))

    ground_truth_dict = generate_dict_from_detection(detection_result_dir)
    
    complete_count = 0
    null_num = 0
    for img_name in ground_truth_dict:
        if not ground_truth_dict[img_name]:
            print("{}: null result.".format(img_name))
            null_num += 1
            continue
        img_file = os.path.join(img_path, img_name)
        if not os.path.exists(img_file):
            print(img_path, " not exist!")
            continue
        doc = Document()
        img = cv2.imread(img_file)
        n_rows, n_cols, channels = img.shape
 
        annotation = doc.createElement('annotation')
        doc.appendChild(annotation)
            
        folder = doc.createElement("folder")
        folder.appendChild(doc.createTextNode("AIA AUTO"))
        annotation.appendChild(folder)
            
        filename = doc.createElement("filename")
        filename.appendChild(doc.createTextNode(img_name))
        annotation.appendChild(filename)
            
        source = doc.createElement("source")
        database = doc.createElement("database")
        database.appendChild(doc.createTextNode("The AUTO Database"))
        anno = doc.createElement("annotation")
        anno.appendChild(doc.createTextNode("HIKDATA_AIA AUTO"))
        image = doc.createElement("image")
        image.appendChild(doc.createTextNode("flickr"))
        source.appendChild(database)
        source.appendChild(anno)
        source.appendChild(image)
        annotation.appendChild(source)
            
        image_size = doc.createElement("size")
        image_width = doc.createElement("width")
        image_width.appendChild(doc.createTextNode(str(n_cols)))
        image_height = doc.createElement("height")
        image_height.appendChild(doc.createTextNode(str(n_rows)))
        image_depth = doc.createElement("depth")
        image_depth.appendChild(doc.createTextNode(str(channels)))
        image_size.appendChild(image_depth)
        image_size.appendChild(image_width)
        image_size.appendChild(image_height)
        annotation.appendChild(image_size)
        
        for box in ground_truth_dict[img_name]:
            # box: [class_name,left, top, right, bottom]
            # 仅仅生成classname中指定的类别

            if box[0] not in labels:
                continue
            # class_name_id = labels[box[0]]
            xmin_v = int(float(box[1]))
            ymin_v = int(float(box[2]))
            xmax_v = int(float(box[3]))
            ymax_v = int(float(box[4]))

            if (xmax_v - xmin_v) > n_cols / 3 or (ymax_v - ymin_v) > n_rows / 3:
                if is_invalid_big_box(box, ground_truth_dict[img_name]):
                    continue
            occl = 0
            
            add_object(doc, annotation, box[0],
                       xmin_v, ymin_v, xmax_v, ymax_v, occl)
        
        xml_name = img_name.split('.')[0]+'.xml'
        xml_file = open(os.path.join(save_xml_path, xml_name), 'w')
        xml_file.write(doc.toprettyxml(indent=' '*4))
        xml_file.close()
        complete_count += 1
        if complete_count % 100 == 0:
            print("Complete {}, Total: {}".format(complete_count, detection_result_num))
    print("*"*80)
    print("Complete {}, No detection result: {} Total: {}".format(
        complete_count, null_num, detection_result_num))


if __name__ == '__main__':
    try:
        ROOT_DIR = sys.argv[1]
        DET_DIR = float(sys.argv[2])
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--root_dir", type=str,
                            default="/home/chujie/PycharmProjects/data_tool/data/voc2012_2",
                            help="VOC dataset root directory")
        PARSER.add_argument("--detection_dir", type=str, default="data/voc2012/result",
                            help="iou threshold")
        ARGS = PARSER.parse_args()
        ROOT_DIR = ARGS.root_dir
        DET_DIR = ARGS.detection_dir
    class_id_name = {"vehicle": 0, "person": 1}
    detection_to_voc_xml(DET_DIR, ROOT_DIR, class_id_name)
